import re
from typing import Any, Dict, Union

from pydantic import BaseModel, ValidationError

from lagent.prompts.parsers.str_parser import StrParser


class CustomFormatParser(StrParser):
    """
    自定义格式解析器
    
    用于解析和格式化XML风格的消息格式，支持字段注释和类型验证
    """

    def _extract_fields_with_metadata(
            self, model: BaseModel) -> Dict[str, Dict[str, Any]]:
        """
        提取模型字段的元数据
        
        参数:
            model: Pydantic模型类
            
        返回:
            字段名到元数据的映射字典，包含:
            - annotation: 字段类型
            - default: 默认值
            - comment: 字段描述
        """
        fields_metadata = {}
        for field_name, field in model.model_fields.items():
            fields_metadata[field_name] = {
                'annotation': field.annotation,
                'default': field.default
                if field.default is not None else '<required>',
                'comment': field.description if field.description else ''
            }
        return fields_metadata

    def format_to_string(self, format_model: BaseModel) -> str:
        """
        将模型格式化为XML风格的字符串
        
        参数:
            format_model: Pydantic模型实例
            
        返回:
            格式化后的XML字符串，包含注释和类型信息
        """
        fields = self._extract_fields_with_metadata(format_model)
        formatted_str = ''
        for field_name, metadata in fields.items():
            # 获取字段注释
            comment = metadata.get('comment', '')
            # 获取字段类型名称
            field_annotation = metadata['annotation'].__name__ if metadata[
                'annotation'] is not None else 'Any'
            # 添加注释（如果有）
            if comment:
                formatted_str += f'<!-- {comment} -->\n'
            # 添加XML标签
            formatted_str += f'<{field_name} type="{field_annotation}">{metadata["default"] if metadata["default"] != "<required>" else ""}</{field_name}>\n'
        return formatted_str

    def parse_response(self, data: str) -> Union[dict, BaseModel]:
        """
        解析XML格式的响应数据
        
        参数:
            data: XML格式的字符串
            
        返回:
            解析后的模型实例
            
        说明:
            支持注释和类型信息的提取，自动选择合适的模型进行验证
        """
        # 匹配XML标签和注释的正则表达式
        pattern = re.compile(r'(<!--\s*(.*?)\s*-->)?\s*<(\w+)[^>]*>(.*?)</\3>',
                           re.DOTALL)
        matches = pattern.findall(data)

        # 构建数据字典
        data_dict = {}
        for _, comment_text, key, value in matches:
            if comment_text:
                self.fields[key]['comment'] = comment_text.strip()
            data_dict[key] = value

        # 选择合适的模型进行验证
        model = self.default_format
        if self.unknown_format and not self._is_valid_format(
                data_dict, self.default_format):
            model = self.unknown_format

        # 验证并返回模型实例
        return model.model_validate(data_dict)

    def _is_valid_format(self, data: Dict, format_model: BaseModel) -> bool:
        """
        验证数据是否符合指定模型格式
        
        参数:
            data: 要验证的数据字典
            format_model: 验证用的模型类
            
        返回:
            bool: 验证是否通过
        """
        try:
            format_model.model_validate(data)
            return True
        except ValidationError:
            return False


# 示例代码
if __name__ == '__main__':
    # 定义默认格式模型
    class DefaultFormat(BaseModel):
        name: str
        age: int

    # 定义未知格式模型
    class UnknownFormat(BaseModel):
        title: str
        year: int

    # 定义模板
    template = """如果了解该问题请按照一下格式回复
                    ```html
                    {format}
                    ```
                    否则请回复
                    ```html
                        {unknown_format}
                        ```
                        """
                        
    # 创建解析器实例
    parser = CustomFormatParser(
        template, 
        default_format=DefaultFormat, 
        unknown_format=UnknownFormat
    )

    # 测试数据
    response = '''
    <!-- User's full name -->
    <name type="str">John Doe</name>
    <!-- User's age -->
    <age type="int">30</age>
    '''

    # 解析并打印结果
    result = parser.parse_response(response)
    print(result)
